{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/MouseLand/cellpose/blob/main/notebooks/train_Cellpose-SAM.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Nb90LCrotIx4"
      },
      "source": [
        "## Cellpose-SAM: superhuman generalization for cellular segmentation\n",
        "\n",
        "Marius Pachitariu, Michael Rariden, Carsen Stringer\n",
        "\n",
        "[paper](https://www.biorxiv.org/content/10.1101/2025.04.28.651001v1) | [code](https://github.com/MouseLand/cellpose)\n",
        "\n",
        "This notebook shows how to process your own 2D or 3D images, saved on Google Drive.\n",
        "\n",
        "This notebook is adapted from the notebook by Pradeep Rajasekhar, inspired by the [ZeroCostDL4Mic notebook series](https://github.com/HenriquesLab/ZeroCostDL4Mic/wiki)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z0s2fz5hUk75"
      },
      "source": [
        "### Make sure you have GPU access enabled by going to Runtime -> Change Runtime Type -> Hardware accelerator and selecting GPU\n",
        "\n",
        "![image.png]()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5-Pn33KNqmE6"
      },
      "source": [
        "### Mount your google drive\n",
        "\n",
        "If you have some images to train on, mount your drive. Alternatively scroll down and download the example images.\n",
        "\n",
        "Run this cell to connect your Google Drive to colab:\n",
        "* Click on the URL.\n",
        "* Sign in your Google Account.\n",
        "\n",
        "You will either have to:\n",
        "* copy the authorisation code and enter it into box below OR\n",
        "* in the new google colab, you can just click \"Allow\" and it should connect."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b8wPXz7PGNLt"
      },
      "outputs": [],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oK-WIju4qmE7"
      },
      "source": [
        "\n",
        "Then click on \"Folder\" icon on the Left, press the refresh button. Your Google Drive folder should now be available here as \"gdrive\".\n",
        "\n",
        "\n",
        "![image.png]()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yH4zTaWwvH9x"
      },
      "source": [
        "Click on the triangle icon and it will allow you to access whole drive. Navigate to the folder containing your images. Once you are there, click on the three dots on the right of the folder and select \"Copy Path\"\n",
        "\n",
        "![image.png]()\n",
        "\n",
        "Copy and paste this path in the **dir** string below"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h_iAN7cAthma"
      },
      "source": [
        "### Install Cellpose-SAM\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "hG3LSmJmLylT"
      },
      "outputs": [],
      "source": [
        "!pip install git+https://www.github.com/mouseland/cellpose.git"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CnKdFgZTqmE9"
      },
      "source": [
        "Check GPU and instantiate model - will download weights."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "5ydQ-fggSiUm"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "from cellpose import models, core, io, plot\n",
        "from pathlib import Path\n",
        "from tqdm import trange\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "io.logger_setup() # run this to get printing of progress\n",
        "\n",
        "#Check if colab notebook instance has GPU access\n",
        "if core.use_gpu()==False:\n",
        "  raise ImportError(\"No GPU access, change your runtime\")\n",
        "\n",
        "model = models.CellposeModel(gpu=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "plEha5EaqmE9"
      },
      "source": [
        "Input directory with your images (if you have them, otherwise use sample images):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-lZP6alpUAfY"
      },
      "outputs": [],
      "source": [
        "# *** change to your google drive folder path ***\n",
        "train_dir = \"/content/gdrive/MyDrive/PATH-TO-FILES/\"\n",
        "if not Path(train_dir).exists():\n",
        "  raise FileNotFoundError(\"directory does not exist\")\n",
        "\n",
        "test_dir = None # optionally you can specify a directory with test files\n",
        "\n",
        "# *** change to your mask extension ***\n",
        "masks_ext = \"_seg.npy\"\n",
        "# ^ assumes images from Cellpose GUI, if labels are tiffs, then \"_masks.tif\"\n",
        "\n",
        "# list all files\n",
        "files = [f for f in Path(train_dir).glob(\"*\") if \"_masks\" not in f.name and \"_flows\" not in f.name and \"_seg\" not in f.name]\n",
        "\n",
        "if(len(files)==0):\n",
        "  raise FileNotFoundError(\"no files found, did you specify the correct folder and extension?\")\n",
        "else:\n",
        "  print(f\"{len(files)} files in folder:\")\n",
        "\n",
        "for f in files:\n",
        "  print(f.name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0JnV_E_OqmE9"
      },
      "source": [
        "### Sample images (optional)\n",
        "\n",
        "You can use our sample images instead of mounting your google drive"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sG96J_V8qmE-"
      },
      "outputs": [],
      "source": [
        "from natsort import natsorted\n",
        "from cellpose import utils\n",
        "from pathlib import Path\n",
        "\n",
        "url = \"https://drive.google.com/uc?id=1HXpLczf7TPCdI1yZY5KV3EkdWzRrgvhQ\"\n",
        "utils.download_url_to_file(url, \"human_in_the_loop.zip\")\n",
        "\n",
        "!unzip human_in_the_loop\n",
        "\n",
        "train_dir = \"human_in_the_loop/train/\"\n",
        "test_dir = \"human_in_the_loop/test/\"\n",
        "\n",
        "masks_ext = \"_seg.npy\"\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dJFJG-mkqmE-"
      },
      "source": [
        "## Train new model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r0umDFliqmE-"
      },
      "outputs": [],
      "source": [
        "from cellpose import train\n",
        "\n",
        "model_name = \"new_model\"\n",
        "\n",
        "# default training params\n",
        "n_epochs = 100\n",
        "learning_rate = 1e-5\n",
        "weight_decay = 0.1\n",
        "batch_size = 1\n",
        "\n",
        "# get files\n",
        "output = io.load_train_test_data(train_dir, test_dir, mask_filter=masks_ext)\n",
        "train_data, train_labels, _, test_data, test_labels, _ = output\n",
        "# (not passing test data into function to speed up training)\n",
        "\n",
        "new_model_path, train_losses, test_losses = train.train_seg(model.net,\n",
        "                                                            train_data=train_data,\n",
        "                                                            train_labels=train_labels,\n",
        "                                                            batch_size=batch_size,\n",
        "                                                            n_epochs=n_epochs,\n",
        "                                                            learning_rate=learning_rate,\n",
        "                                                            weight_decay=weight_decay,\n",
        "                                                            nimg_per_epoch=max(2, len(train_data)), # can change this\n",
        "                                                            model_name=model_name)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gj0EdXtcqmE-"
      },
      "source": [
        "## Evaluate on test data (optional)\n",
        "\n",
        "If you have test data, check performance"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y2Gv4KnSqmE-"
      },
      "outputs": [],
      "source": [
        "from cellpose import metrics\n",
        "\n",
        "model = models.CellposeModel(gpu=True,\n",
        "                             pretrained_model=new_model_path)\n",
        "\n",
        "# run model on test images\n",
        "masks = model.eval(test_data, batch_size=32)[0]\n",
        "\n",
        "# check performance using ground truth labels\n",
        "ap = metrics.average_precision(test_labels, masks)[0]\n",
        "print('')\n",
        "print(f'>>> average precision at iou threshold 0.5 = {ap[:,0].mean():.3f}')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OddRFdtEqmE-"
      },
      "source": [
        "plot masks"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9MUrvy5JqmE-"
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(12,8), dpi=150)\n",
        "for k,im in enumerate(test_data):\n",
        "    img = im.copy()\n",
        "    plt.subplot(3,len(test_data), k+1)\n",
        "    img = np.vstack((img, np.zeros_like(img)[:1]))\n",
        "    img = img.transpose(1,2,0)\n",
        "    plt.imshow(img)\n",
        "    plt.axis('off')\n",
        "    if k==0:\n",
        "        plt.title('image')\n",
        "\n",
        "    plt.subplot(3,len(test_data), len(test_data) + k+1)\n",
        "    plt.imshow(masks[k])\n",
        "    plt.axis('off')\n",
        "    if k==0:\n",
        "        plt.title('predicted labels')\n",
        "\n",
        "    plt.subplot(3,len(test_data), 2*len(test_data) + k+1)\n",
        "    plt.imshow(test_labels[k])\n",
        "    plt.axis('off')\n",
        "    if k==0:\n",
        "        plt.title('true labels')\n",
        "plt.tight_layout()"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "cellpose",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
